# File: util/symbolic_manager.py
from typing import Optional

import numpy as np

class SymbolicManager:
    """Manages the vocabulary and conversion between symbolic tokens and integer indices."""

    def __init__(self, num_processes: int, num_resources: int):
        self.num_processes = num_processes
        self.num_resources = num_resources

        # --- 1. Define all special and structural tokens ---
        self.special_tokens = [
            '[PAD]', '[START]', '[END]', '[DECIDE]',
            '[WAIT_FOR]', '[HELD_BY]', '[CYCLE_FOUND]', '[END_CYCLE]',
            '[EVAL_COST]', '[PREEMPT]', '[ROLLBACK]', '[NO_OP]',
            '[STATE_SAFE]'
        ]
        self.process_tokens = [f'P{i}' for i in range(num_processes)]
        self.resource_tokens = [f'R{i}' for i in range(num_resources)]

        # --- 2. Build the complete vocabulary and token-to-index mapping ---
        self.vocab = self.special_tokens + self.process_tokens + self.resource_tokens
        self.token_to_idx = {token: i for i, token in enumerate(self.vocab)}

        # --- 3. (THE FIX) Create the reverse index-to-token mapping ---
        self.idx_to_token = {i: token for token, i in self.token_to_idx.items()}  # <-- THIS IS THE NEW LINE

        # --- 4. Store key token indices as attributes for easy access ---
        self.PAD = self.token_to_idx['[PAD]']
        self.START = self.token_to_idx['[START]']
        self.END = self.token_to_idx['[END]']
        self.DECIDE = self.token_to_idx['[DECIDE]']
        self.CYCLE_FOUND = self.token_to_idx['[CYCLE_FOUND]']
        self.EVAL_COST = self.token_to_idx['[EVAL_COST]']
        self.NO_OP = self.token_to_idx['[NO_OP]']

    def get_vocab_size(self) -> int:
        """Returns the total size of the vocabulary."""
        return len(self.vocab)

    def parse_action_from_trace(self, trace_tokens: np.ndarray) -> Optional[dict]:
        """Parses the final decision from a trace."""
        try:
            tokens = [self.idx_to_token.get(t) for t in trace_tokens]
            if '[DECIDE]' in tokens:
                idx = tokens.index('[DECIDE]')
                action_token = tokens[idx + 1]
                process_token = tokens[idx + 2]

                action_type = None
                if action_token == '[PREEMPT]':
                    action_type = 'preempt'
                elif action_token == '[ROLLBACK]':
                    action_type = 'rollback'

                if action_type and process_token.startswith('P'):
                    process_id = int(process_token[1:])
                    return {'type': action_type, 'process_id': process_id}
        except (ValueError, IndexError, TypeError):
            # If parsing fails, return None
            return None
        return None

    def parse_trace_to_assertions(self, trace_tokens: np.ndarray) -> list:
        """
        Parses a sequence of tokens into a list of structured assertion dictionaries.
        This version is updated to also parse the final [DECIDE] action.
        """
        assertions = []
        i = 0
        tokens = [self.idx_to_token.get(t) for t in trace_tokens]

        while i < len(tokens):
            symbol = tokens[i]
            if symbol is None:
                i += 1
                continue

            try:
                # --- [新增逻辑] 解析 DECIDE 断言 ---
                if symbol == '[DECIDE]' and i + 2 < len(tokens):
                    action_symbol = tokens[i + 1]
                    process_symbol = tokens[i + 2]

                    action_type = None
                    if action_symbol == '[PREEMPT]':
                        action_type = 'preempt'
                    elif action_symbol == '[ROLLBACK]':
                        action_type = 'rollback'

                    if action_type and process_symbol and process_symbol.startswith('P'):
                        p_id = int(process_symbol[1:])
                        assertions.append({'type': 'DECIDE', 'action': action_type, 'process': p_id})
                        i += 3  # 跳过 [DECIDE], [ACTION], [PROCESS]
                        continue
                # --- 新增逻辑结束 ---

                elif symbol == '[WAIT_FOR]' and i > 0 and i + 1 < len(tokens):
                    prev_symbol = tokens[i - 1]
                    next_symbol = tokens[i + 1]
                    if prev_symbol and prev_symbol.startswith('P') and next_symbol and next_symbol.startswith('R'):
                        p_id = int(prev_symbol[1:])
                        r_id = int(next_symbol[1:])
                        assertions.append({'type': 'WAIT', 'process': p_id, 'resource': r_id})
                        i += 2
                        continue

                elif symbol == '[HELD_BY]' and i > 0 and i + 1 < len(tokens):
                    prev_symbol = tokens[i - 1]
                    next_symbol = tokens[i + 1]
                    if prev_symbol and prev_symbol.startswith('R') and next_symbol and next_symbol.startswith('P'):
                        r_id = int(prev_symbol[1:])
                        p_id = int(next_symbol[1:])
                        assertions.append({'type': 'HELD_BY', 'resource': r_id, 'process': p_id})
                        i += 2
                        continue

                elif symbol == '[CYCLE_FOUND]' and i + 1 < len(tokens):
                    path = []
                    j = i + 1
                    while j < len(tokens) and tokens[j] != '[END_CYCLE]':
                        cycle_symbol = tokens[j]
                        if cycle_symbol and cycle_symbol.startswith('P'):
                            path.append(int(cycle_symbol[1:]))
                        j += 1

                    if j < len(tokens):  # 确保找到了 [END_CYCLE]
                        assertions.append({'type': 'CYCLE', 'path': path})
                        i = j + 1
                        continue

            except (ValueError, IndexError, TypeError):
                pass

            i += 1
        return assertions

